import numpy as np
from sfepy.discrete.fem import Mesh, FEDomain, Field
from sfepy.discrete import Problem, Functions, Equations, Region
from sfepy.base.base import assert_, output, Struct
from sfepy.solvers import Solver
from sfepy.mechanics.matcoefs import stiffness_from_lame
from sfepy.discrete import (FieldVariable, Material, Integral, Integrals,
                            Equation, Equations, Problem)
from sfepy.mechanics.matcoefs import stiffness_from_youngpoisson
from sfepy.terms import Term
from sfepy.discrete.conditions import Conditions, EssentialBC
import scipy.sparse.linalg as sla
import matplotlib.pyplot as plt
import pandas as pd
import os
import structAIDetect.simdatagen.crackgenerate as crackgenerate
import itertools
from structAIDetect.simdatagen.modaldatagen import modal_analysis, cellinfo_generate
import argparse


N_EIGS = 5

def main():
    parser = argparse.ArgumentParser()
    parser.add_argument('--filename', '-F', help='输入几何建模的文件名（包括扩展名）')
    args = parser.parse_args()
    filename = args.filename

    # 1. 加载网格
    modelgeo_path = os.path.join(os.getcwd(), "data", "geo")

    output_dir = os.path.join(os.getcwd(), 'data', 'sim')

    dim = 3
    axis = 1
    axis_fix = 0  # 沿x轴方向固定

    mesh = Mesh.from_file(os.path.join(modelgeo_path, filename))
    domain = FEDomain('domain', mesh)
    bbox = domain.get_mesh_bounding_box()
    min_coor, max_coor = bbox[:, axis]
    min_side_coor, max_side_coor = bbox[:, axis_fix]
    eps = 1e-8 * (max_coor - min_coor)
    eps_side = 1e-8 * (max_side_coor - min_side_coor)
    ax = 'xyz'[:dim][axis]
    aside = 'xyz'[:dim][axis_fix]

    omega = domain.create_region('Omega', 'all')
    left_side = domain.create_region('Rightside',
                                    'vertices in (%s < %.10f)'
                                    % (aside, min_side_coor + eps_side),
                                    'facet')
    right_side = domain.create_region('Rightside',
                                    'vertices in (%s > %.10f)'
                                    % (aside, max_side_coor - eps_side),
                                    'facet')

    # 生成裂纹位置
    # cavity_list = []
    crack_y_list = np.arange(np.abs(4/6 * (max_coor - min_coor)) // 6) * 6 + (5/6 * min_coor + 1/6 * max_coor)
    crack_x_list = np.arange(np.abs(1/2 * (max_side_coor - min_side_coor)) // 6) * 6 + (3/4 * min_side_coor + 1/4 * max_side_coor)
    crack_dir_list = np.array([-1, 1])
    combinations = list(itertools.product(crack_y_list, crack_x_list, crack_dir_list))
    target_corr_z = 0.1
    sim_idx = 0
    for target_corr_y, target_corr_x, crack_dir in combinations:
        target_corr = np.array([target_corr_x, target_corr_y, target_corr_z])
        if crack_dir > 0:
            cavity = crackgenerate.get_cavity_region(domain, omega, 
                                                    target_corr, right_side, crack_dir)
        else:
            cavity = crackgenerate.get_cavity_region(domain, omega, 
                                                    target_corr, left_side, crack_dir)
        # cavity_list.append(cavity)

        mesh = Mesh.from_file(os.path.join(modelgeo_path, filename))
        domain = FEDomain('domain', mesh)
        bbox = domain.get_mesh_bounding_box()
        min_coor, max_coor = bbox[:, axis]
        min_side_coor, max_side_coor = bbox[:, axis_fix]
        eps = 1e-8 * (max_coor - min_coor)
        eps_side = 1e-8 * (max_side_coor - min_side_coor)
        ax = 'xyz'[:dim][axis]
        aside = 'xyz'[:dim][axis_fix]
        omega = domain.create_region('Omega', 'all')
        # 设置约束区域
        fix_bottom = domain.create_region('Bottom',
                                    'vertices in (%s < %.10f)'
                                    % (ax, min_coor + eps),
                                    'facet')

        # 2. 定义材料属性（钢为例）
        youngs_modulus = 210e3  # MPa (mm 量纲)
        poissons_ratio = 0.3
        density = 7.85e-6       # kg/mm³

        # 3. 定义求解器
        eig_conf = Struct(name='evp', kind='eig.scipy', method='eig')

        # 4. 定义有限元场（六面体线性单元）
        order = 1
        field = Field.from_args('fu', np.float64, 'vector', omega,
                                    approx_order=order)
        u = FieldVariable('u', 'unknown', field)
        v = FieldVariable('v', 'test', field, primary_var_name='u')

        # 5. 定义材料
        mtx_d = stiffness_from_youngpoisson(dim, young=youngs_modulus, poisson=poissons_ratio)
        m = Material('m', D=mtx_d, rho=density)
        integral = Integral('i', order=2*order)
        t1 = Term.new('dw_lin_elastic(m.D, v, u)', integral, omega, m=m, v=v, u=u)
        t2 = Term.new('dw_dot(m.rho, v, u)', integral, omega, m=m, v=v, u=u)
        eq1 = Equation('stiffness', t1)
        eq2 = Equation('mass', t2)
        lhs_eqs = Equations([eq1, eq2])
        pb = Problem('modal', equations=lhs_eqs)

        # 6. 定义边界条件
        fixed = EssentialBC('Fixed', fix_bottom, {'u.all' : 0.0})
        pb.time_update(ebcs=Conditions([fixed]))
        pb.update_materials()
        eig_solver = Solver.any_from_conf(eig_conf)
        freqs, vecs = modal_analysis(eq1, eq2, pb, cavity, N_EIGS, eig_solver)
        cellinfo, cellstate = cellinfo_generate(domain, omega, vecs, cavity)
        # shape_plot(domain, vecs, 5, freqs)
        np.savez(os.path.join(output_dir, filename[:-4] + 'sim' + str(sim_idx) + '.npz'), 
                cellinfo=cellinfo, cellstate=cellstate, freqs=freqs)
        sim_idx += 1



if __name__ == '__main__':
    main()